// Copyright 2022 Google LLC
//
// This source code is licensed under the BSD-style license found in the
// LICENSE file in the root directory of this source tree.

#include <xnnpack/assembly.h>

.syntax unified

// void xnn_cs16_bfly4_samples1_ukernel__asm_aarch32_neon_x4(
//     size_t batch,                         r0
//     size_t samples,                       (unused)
//     int16_t* data,                        r2
//     const int16_t* twiddle,               (unused)
//     size_t stride)                        (unused)

// d8-d15, r12-r11,r14(lr) need to be preserved if used. r13(sp),r15(pc) are reserved.

// Register usage
// vout0 r2 q0
// vout1    q1
// vout2    q2
// vout3    q3
// vtmp3    q8
// vtmp4    q9
// vtmp5    q10
// vtmp0    q11

// vdiv4    q12
// vnegr    q13

BEGIN_FUNCTION xnn_cs16_bfly4_samples1_ukernel__asm_aarch32_neon_x4
        .arm
#ifndef __APPLE__
        .arch   armv7-a
        .fpu    neon
#endif
        SUBS    r0, r0, 4               // batch
        VMVN.U16 q12, 57344             // 8191
        VMOV.I32 q13, 0x0001ffff        // vnegr
        BLO     1f

        MOV     r3, r2                  // output = input for post inc

        // batch of 4 main loop
0:
        VLD4.32 {d0,d2,d4,d6}, [r2]!    // input first 2 batch
        VLD4.32 {d1,d3,d5,d7}, [r2]!    // input second 2 batch
        SUBS    r0, r0, 4               // batch
        VQRDMULH.S16 q1, q1, q12        // vout1 /= 4
        VQRDMULH.S16 q3, q3, q12        // vout3 /= 4
        VQRDMULH.S16 q0, q0, q12        // vout0 /= 4
        VQRDMULH.S16 q2, q2, q12        // vout2 /= 4

        VSUB.I16 q9, q1, q3             // vtmp4 = vout1 - vout3
        VADD.I16 q8, q1, q3             // vtmp3 = vout1 + vout3

        VMUL.S16 q9, q9, q13            // vrev4 = vtmp4 -r, i
        VADD.I16 q11, q0, q2            // vtmp0 = vout0 + vout2
        VSUB.I16 q10, q0, q2            // vtmp5 = vout0 - vout2

        VADD.I16 q0, q11, q8            // vout0 = vtmp0 + vtmp3
        VSUB.I16 q2, q11, q8            // vout2 = vtmp0 - vtmp3

        VREV32.16 q9, q9                // vrev4 = vtmp4 i, -r
        VADD.I16 q1, q10, q9            // vout1 = vtmp5 + vrev4
        VSUB.I16 q3, q10, q9            // vout3 = vtmp5 - vrev4

        VST4.32 {d0,d2,d4,d6}, [r3]!    // output first 2 batch
        VST4.32 {d1,d3,d5,d7}, [r3]!    // output second 2 batch

        BHS     0b

1:
        ANDS    r0, r0, 3               // batch remainder?
        BXEQ    lr

        // Remainder batch of 1 to 3
2:
        VLD4.32 {d0[0],d1[0],d2[0],d3[0]}, [r2]  // input 1 batch
        SUBS    r0, r0, 1               // batch
        VQRDMULH.S16 d1, d1, d24        // vout1 /= 4
        VQRDMULH.S16 d3, d3, d24        // vout3 /= 4
        VQRDMULH.S16 d0, d0, d24        // vout0 /= 4
        VQRDMULH.S16 d2, d2, d24        // vout2 /= 4

        VSUB.I16 d5, d1, d3             // vtmp4 = vout1 - vout3
        VADD.I16 d4, d1, d3             // vtmp3 = vout1 + vout3
        VMUL.S16 d5, d5, d26            // vrev4 = vtmp4 -r, i

        VADD.I16 d7, d0, d2             // vtmp0 = vout0 + vout2
        VSUB.I16 d6, d0, d2             // vtmp5 = vout0 - vout2

        VADD.I16 d0, d7, d4             // vout0 = vtmp0 + vtmp3
        VSUB.I16 d2, d7, d4             // vout2 = vtmp0 - vtmp3

        VREV32.16 d5, d5                // vrev4 = vtmp4 i, -r
        VADD.I16 d1, d6, d5             // vout1 = vtmp5 + vrev4
        VSUB.I16 d3, d6, d5             // vout3 = vtmp5 - vrev4

        VST4.32 {d0[0],d1[0],d2[0],d3[0]}, [r2]! // output 1 batch
        BHI     2b

        BX      lr

END_FUNCTION xnn_cs16_bfly4_samples1_ukernel__asm_aarch32_neon_x4

#ifdef __ELF__
.section ".note.GNU-stack","",%progbits
#endif
